package cn.showcon.firstapp.service;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.OnnxEmbeddingModel;
import dev.langchain4j.model.embedding.onnx.PoolingMode;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.store.embedding.CosineSimilarity;

import java.util.HashMap;
import java.util.Map;


/**
 * @Author Xue Lanbin
 */
public class InProcessEmbeddingModelExamples {

    static class LocalCustomModelExample1 {

        public static void main(String[] args) {
            String text = "Let's demonstrate that embedding can be done within a Java process and entirely offline.";

            String pathToModel = "D:\\xuelb\\llm\\resources\\model\\embedding\\shibing624_text2vec-base-chinese\\onnx\\model.onnx";
            String pathToTokenizer = "D:\\xuelb\\llm\\resources\\model\\embedding\\shibing624_text2vec-base-chinese\\onnx\\tokenizer.json";
            PoolingMode poolingMode = PoolingMode.MEAN;
            OnnxEmbeddingModel localEmbeddingModel = new OnnxEmbeddingModel(pathToModel, pathToTokenizer, poolingMode);
            System.out.println("localEmbeddingModel.dimension(): " + localEmbeddingModel.dimension());

            Embedding inProcessEmbedding = localEmbeddingModel.embed(text).content();
            System.out.println(inProcessEmbedding);
            System.out.println("Dimension: " + inProcessEmbedding.dimension());
            System.out.println("Vector Size: " + inProcessEmbedding.vectorAsList().size());


            Embedding embedding1 = localEmbeddingModel.embed("Embedding").content();
            Embedding embedding2 = localEmbeddingModel.embed("Embedding").content();
            double similarity = CosineSimilarity.between(embedding1, embedding2);
            System.out.println(similarity);

        }
    }

    static class AllMinilmL6V2Example1 {

        public static void main(String[] args) {

            // requires "langchain4j-embeddings-all-minilm-l6-v2" Maven/Gradle dependency, see pom.xml
            // https://huggingface.co/optimum/all-MiniLM-L6-v2/tree/main
            EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

            Embedding embedding1 = embeddingModel.embed("huggingFaceEmbedding").content();
            System.out.println("embeddingModel.dimension(): " + embeddingModel.dimension());

            Embedding embedding2 = embeddingModel.embed("localEmbedding").content();
            double similarity = CosineSimilarity.between(embedding1, embedding2);
            System.out.println(similarity);

            Map<String, Object> variables = new HashMap<>();
            variables.put("name", "alex");

            PromptTemplate promptTemplate = PromptTemplate.from("Name: {{name}}");
            Prompt prompt = promptTemplate.apply(variables);

            System.out.println(prompt);
            System.out.println(prompt.text());
            System.out.println(prompt.toSystemMessage());
            System.out.println(prompt.toUserMessage());


        }
    }
}
